# -*- coding: utf-8 -*-
from __future__ import print_function

import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import torch.optim as optim

from ddl_platform.ddlib import job
from models import FCN5Net


class MnistJob(job.Job):
    def build_dataset(self):
        config = self.config()['dataset']

        data_dir = config['data_dir']

        trans = []
        trans.extend([transforms.ToTensor(),
                      transforms.Normalize((0.1307,), (0.3081,))
                        ])
        trainset = torchvision.datasets.MNIST(data_dir, train=True, download=True,
                    transform=transforms.Compose(trans))

        testset = torchvision.datasets.MNIST(data_dir, train=False, transform=transforms.Compose(trans))
        return trainset, testset


    def build_model(self):
        return FCN5Net()

    def build_optimizer(self):
        model = self.model()
        config = self.config()['optimizer']
        optimizer = optim.SGD(model.parameters(), 
                lr=config['lr'],
                momentum=config['momentum'])
        return optimizer

    def build_criterion(self):
        return nn.CrossEntropyLoss()

    def cal_eval_performance(self, batch_outputs, batch_inputs):
        return super().cal_eval_performance(batch_outputs, batch_inputs)


#def train():
#    conf_yaml = 'mnist.yaml'
#    job = MnistJob(conf_yaml)
#    t = trainer.Trainer(job)
#    t.fit()
#
#if __name__ == '__main__':
#    train()
